from transformers import AutoTokenizer, AutoModelForCausalLM

# 选择LLaMA模型
model_name = 'meta-llama/Llama-2-7b'  # 可以选择 'meta-llama/LLaMA-13B'、'meta-llama/LLaMA-30B' 等

# 加载LLaMA的tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained('E:/MyCode/PyCharmWorkSpace/ManipLLM/test/ckpts/llama_model_weights')

# 输入文本
input_text = "这是一个示例文本，用于计算LLaMA模型的Token数量。"

# 使用tokenizer对输入文本进行编码
encoded_input = tokenizer(input_text, return_tensors='pt')

# 计算输入文本的Token数量
input_token_count = len(encoded_input['input_ids'][0])



# 打印Token数量
print(f"输入文本的Token数量: {input_token_count}")
